
import chisel3._
import chisel3.util._

class InstFetch extends Module {
  val io = IO(new Bundle {
    val stall   = Input(Bool())
    val br_en   = Input(Bool())
    val br_addr = Input(UInt(32.W))
    val pc      = Output(UInt(32.W))
    val inst    = Output(UInt(32.W)) 
    val imem_hs = Output(Bool()) 
    val imem    = new CacheBusIO
  })

//  val if_axi_id = 1.U(AxiParameters.AxiIdWidth.W) 

  val s_init :: s_idle :: s_req :: s_wait :: Nil = Enum(4)
  val state = RegInit(s_init)

  val req = io.imem.req
  val resp = io.imem.resp
  val stall = io.stall

  val pc = RegInit("h80000000".U(32.W))
  val pc_nxt = Mux(io.br_en, io.br_addr, pc + 4.U)
  val inst = RegInit(0.U(32.W))

//  req.bits.id := if_axi_id
  req.bits.addr := pc.asUInt()
  req.bits.ren := true.B          // read-only imem
  req.bits.wdata := 0.U
  req.bits.wmask := 0.U
  req.bits.wen := false.B
  req.bits.size := "b11".U
  req.valid := (state === s_req) && !stall 

  resp.ready := (state === s_wait)

  //val resp_success = resp.fire() && resp.bits.rlast && (resp.bits.id === if_axi_id)

  switch (state) {
    is (s_init) {
      state := s_req
    }
    is (s_idle) {
      pc := Mux(stall, pc, pc_nxt)
      state := Mux(stall, s_idle, s_req)
    }
    is (s_req) {
      when (!stall && req.fire()) {
        state := s_wait
      }
    }
    is (s_wait) {
      when (resp.fire()) {
          inst := Mux(pc(2), resp.bits.rdata(63, 32), resp.bits.rdata(31, 0))
          state := s_idle
      }
    }
  }

  io.pc := Mux(state === s_idle, pc, 0.U)
  io.inst := Mux(state === s_idle, inst, 0.U)
  io.imem_hs := resp.valid && resp.ready
}

